# train.py
import numpy as np
import torch.nn as nn
import torch
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.pyplot as plt

# 导入自定义模块
from utils.data_processing import load_data, preprocess_data
from models.model import 丰度预测模型

#   超参数
BATCH_SIZE = 2  # 样本少，批次设小
EPOCHS = 500  # 训练轮次
LR = 0.001  # 学习率
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 用GPU加速（可选）


def train():
    # 1. 加载并预处理数据
    s, p = load_data()
    s_train, s_test, p_train, p_test = preprocess_data(s, p)

    # 转换为PyTorch张量
    s_train = torch.tensor(s_train, dtype=torch.float32).to(DEVICE)
    p_train = torch.tensor(p_train, dtype=torch.float32).to(DEVICE)
    s_test = torch.tensor(s_test, dtype=torch.float32).to(DEVICE)
    p_test = torch.tensor(p_test, dtype=torch.float32).to(DEVICE)

    # 创建DataLoader（批量加载数据）
    train_dataset = TensorDataset(s_train, p_train)
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

    # 2. 初始化模型、损失函数和优化器
    model = 丰度预测模型().to(DEVICE)
    criterion = nn.MSELoss()  # 均方误差（回归任务常用）
    optimizer = optim.Adam(model.parameters(), lr=LR)

    # 3. 训练模型
    train_losses = []
    test_losses = []

    for epoch in range(EPOCHS):
        model.train()  # 训练模式
        train_loss = 0.0

        for batch_s, batch_p in train_loader:
            optimizer.zero_grad()  # 清零梯度
            outputs = model(batch_s)  # 预测丰度
            loss = criterion(outputs, batch_p)  # 计算损失
            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数
            train_loss += loss.item() * batch_s.size(0)

        # 计算平均训练损失
        train_loss /= len(train_loader.dataset)
        train_losses.append(train_loss)

        # 测试集评估（不更新参数）
        model.eval()
        with torch.no_grad():
            test_outputs = model(s_test)
            test_loss = criterion(test_outputs, p_test).item()
            test_losses.append(test_loss)

        # 打印训练 进度
        if (epoch + 1) % 50 == 0:
            print(f"Epoch [{epoch + 1}/{EPOCHS}], Train Loss: {train_loss:.6f}, Test Loss: {test_loss:.6f}")

    # 4. 保存模型
    torch.save(model.state_dict(), "models/丰度预测模型.pth")
    print("模型已保存到 models/丰度预测模型.pth")

    # 5. 绘制损失曲线
    plt.plot(range(1, EPOCHS + 1), train_losses, label="Train Loss")
    plt.plot(range(1, EPOCHS + 1), test_losses, label="Test Loss")
    plt.xlabel("Epoch")
    plt.ylabel("MSE Loss")
    plt.title("Training and Testing Loss")
    plt.legend()
    plt.savefig("loss_curve.png")
    plt.show()


if __name__ == "__main__":
    train()